from ultralytics import YOLO
import cv2
from .logger import get_logger

class Detector:
    def __init__(self, cfg):
        self.log = get_logger("Detector")
        self.model = YOLO(cfg["weights"])
        self.conf = cfg["conf"]
        self.target = set(cfg["target_classes"])

    def __call__(self, frame):
        self.log.debug("开始推理")
        results = self.model(frame, verbose=False)[0]
        boxes = []
        for box in results.boxes:
            cls_name = results.names[int(box.cls)]
            if cls_name in self.target and float(box.conf) >= self.conf:
                boxes.append({
                    "cls": cls_name,
                    "conf": float(box.conf),
                    "xyxy": box.xyxy.cpu().numpy().astype(int).tolist()
                })
        self.log.info("检出目标 %d 个", len(boxes))
        annotated = results.plot() if boxes else frame
        return boxes, annotated